import numpy as np
# from scene import Scene
from collections import deque
import vtkplotlib as vpl


class AABB:
    """
    Класс, описывающий ограничивающий ограничивающий параллепипед (ААВВ)
    """
    def __init__(self, botleft, topright) -> None:
        self.botleft = botleft
        self.topright = topright

    def create_box(self):
        self.vertex = np.zeros((8,3), dtype=np.float32) #массив с вершинами
        self.vertex[0] = self.botleft
        self.vertex[1] = [self.botleft[0], self.topright[1], self.botleft[2]]
        self.vertex[2] = [self.topright[0], self.topright[1], self.botleft[2]]
        self.vertex[3] = [self.topright[0], self.botleft[1], self.botleft[2]]
        self.vertex[4] = [self.botleft[0], self.botleft[1], self.topright[2]]
        self.vertex[5] = [self.botleft[0], self.topright[1], self.topright[2]]
        self.vertex[6] = self.topright
        self.vertex[7] = [self.topright[0], self.botleft[1], self.topright[2]]

        """треугольники параллепипеда"""
        self.triangles_aabb = np.array([[self.vertex[0], self.vertex[4], self.vertex[1]],
                              [self.vertex[1], self.vertex[4], self.vertex[5]],
                              [self.vertex[4], self.vertex[7], self.vertex[5]],
                              [self.vertex[5], self.vertex[7], self.vertex[6]],
                              [self.vertex[3], self.vertex[7], self.vertex[6]],
                              [self.vertex[2], self.vertex[3], self.vertex[6]],
                              [self.vertex[0], self.vertex[2], self.vertex[3]],
                              [self.vertex[0], self.vertex[1], self.vertex[2]],
                              [self.vertex[1], self.vertex[6], self.vertex[2]],
                              [self.vertex[1], self.vertex[5], self.vertex[6]],
                              [self.vertex[0], self.vertex[7], self.vertex[4]],
                              [self.vertex[0], self.vertex[3], self.vertex[7]]])

        self.face = np.array([[0, 4, 1],
                              [1, 4, 5],
                              [4, 7, 5],
                              [5, 7, 6],
                              [3, 7, 6],
                              [2, 3, 6],
                              [0, 2, 3],
                              [0, 1, 2],
                              [1, 6, 2],
                              [1, 6, 6],
                              [0, 7, 4],
                              [0, 3, 7]])


class BVHNode:
    """класс, описывающий узел"""
    def __init__(self, aabb, triangles, child1, child2):
        self.aabb = aabb
        self.triangles = triangles
        self.child1 = child1
        self.child2 = child2


class SerialBVHNode:
    """
    aabb      : see AABB class avobe
    tristart  : first triangle index of this node
    ntris     : number of triangles of this node
    child1, 2 : child node index, [Note]-1 means there is no child
    """
    def __init__(self, aabb, tristart, ntris):
        self.aabb = aabb
        self.tristart = tristart
        self.ntris = ntris
        self.child1 = -1
        self.child2 = -1
        self.parent = -1


def build_bvh(triangles, min_tri, mode='sah'):
    """
    построение дерева
    :param min_tri:
    :param mode:
    :param triangles: массив полигонов
    :return:
    """
    aabb = AABB(triangles.min(axis=1).min(axis=0), triangles.max(axis=1).max(axis=0))
    if triangles.shape[0] <= min_tri:  #условие для остановки построения
        return BVHNode(aabb, triangles, None, None)

    else:
        if mode == 'mean':
            centroids = np.mean(triangles, axis=1)
            sortind = np.argmax(aabb.topright - aabb.botleft)
            split = (aabb.topright[sortind] + aabb.botleft[sortind])/2
            sorttris = triangles[np.argsort(centroids[:, sortind])]
            lefttris = sorttris[centroids[:, sortind] < split]
            righttris = sorttris[centroids[:, sortind] >= split]
            if len(lefttris) == 0 or len(righttris) == 0:
                return BVHNode(aabb, triangles, None, None)
            node1 = build_bvh(lefttris, mode='mean')
            node2 = build_bvh(righttris, mode='mean')
            return BVHNode(aabb, np.zeros((0, 3, 3), dtype=np.float32), node1, node2)

        if mode == 'median':
            centroids = np.mean(triangles, axis=1)
            sortind = np.argmax(aabb.topright - aabb.botleft)
            sorttris = triangles[np.argsort(centroids[:, sortind])]
            medind = len(sorttris) // 2
            lefttris = sorttris[0:medind]
            righttris = sorttris[medind:]
            node1 = build_bvh(lefttris, min_tri, mode='median')
            node2 = build_bvh(righttris, min_tri, mode='median')
            return BVHNode(aabb, np.zeros((0, 3, 3), dtype=np.float32), node1, node2)

        if mode == 'sah':
            mincost = np.inf
            minaxis = -1
            minsplit = -1
            centroids = np.mean(triangles, axis=1)
            for axis in range(3):
                sortedtris = triangles[np.argsort(centroids[:, axis])]
                N = len(sortedtris)//30
                c = len(sortedtris[::N])
                leftaabb = AABB(np.ones((3,), dtype=np.float32) * np.inf,
                                np.ones((3,), dtype=np.float32) * -np.inf)
                leftcost = np.zeros((c,), dtype=np.float32)
                for i, tri in enumerate(sortedtris[::N]):
                    leftaabb.botleft = np.minimum(leftaabb.botleft, tri.min(axis=0))
                    leftaabb.topright = np.maximum(leftaabb.topright, tri.max(axis=0))
                    diag = np.abs(leftaabb.topright - leftaabb.botleft)
                    left_index = i
                    leftcost[left_index] = (diag[0] * diag[1] + diag[1] * diag[2] + diag[2] * diag[0]) * i

                rightaabb = AABB(np.ones((3,), dtype=np.float32) * np.inf,
                                 np.ones((3,), dtype=np.float32) * -np.inf)
                rightcost = np.zeros((c,), dtype=np.float32)
                for i, tri in enumerate(sortedtris[::-N]):
                    rightaabb.botleft = np.minimum(rightaabb.botleft, tri.min(axis=0))
                    rightaabb.topright = np.maximum(rightaabb.topright, tri.max(axis=0))
                    diag = np.abs(rightaabb.topright - rightaabb.botleft)
                    right_index = c - 1 - i
                    rightcost[right_index] = (diag[0] * diag[1] + diag[1] * diag[2] + diag[2] * diag[0]) * i
                if ((leftcost + rightcost).min() < mincost):
                    mincost = (leftcost + rightcost).min()
                    minaxis = axis
                    minsplit = np.argmin(leftcost + rightcost)

            sortedindices = np.argsort(centroids[:, minaxis])
            sortedtris = triangles[sortedindices]
            if (minsplit*N) == 0:
                return BVHNode(aabb, triangles, None, None)
            node1 = build_bvh(sortedtris[:minsplit*N], min_tri, mode='sah')
            node2 = build_bvh(sortedtris[minsplit*N:], min_tri, mode='sah')
            return BVHNode(aabb, np.zeros((0, 3, 3), dtype=np.float32), node1, node2)


        if mode == 'sah1':
            sah = np.inf
            axis_sah = -1
            N = 20
            diag = np.abs(aabb.topright - aabb.botleft)
            sparent = diag[0]*diag[1]+diag[1]*diag[2]+diag[0]*diag[2]

            centroids = np.mean(triangles, axis=1)
            for axis in range(3):
                # sah_planes = np.ones()
                w = (aabb.topright[axis] + aabb.botleft[axis]) / 2
                split_planes = np.array([w*i/(N+1) for i in range(1, (N+1))])
                # print(split_planes)
                sah_planes = np.ones((len(split_planes),), dtype=np.float32)*np.inf
                sorttris = triangles[np.argsort(centroids[:, axis])]
                # sortedtris = triangles[np.argsort(centroids[:, axis])]
                # N = len(sortedtris) // 10
                # c = len(sortedtris[::N])2) * -np.inf)

                for i, plane in enumerate(split_planes):
                    triangles_left = sorttris[centroids[:, axis] < plane]
                    n_triangles_left = len(triangles_left)
                    triangles_right = sorttris[centroids[:, axis] >= plane]
                    n_triangles_right = len(triangles_right)
                    if n_triangles_left == 0:
                        aabbright = AABB(triangles_right.min(axis=1).min(axis=0),
                                         triangles_right.max(axis=1).max(axis=0))
                        diagright = np.abs(aabbright.topright - aabbright.botleft)
                        sright = diagright[0] * diagright[1] + diagright[1] * diagright[2] + diagright[0] * diagright[2]
                        sah_plane = 1 + ((n_triangles_right * sright) / sparent)
                        sah_planes[i] = sah_plane
                    elif n_triangles_right == 0:
                        aabbleft = AABB(triangles_left.min(axis=1).min(axis=0), triangles_left.max(axis=1).max(axis=0))
                        diagleft = np.abs(aabbleft.topright - aabbleft.botleft)
                        sleft = diagleft[0] * diagleft[1] + diagleft[1] * diagleft[2] + diagleft[0] * diagleft[2]
                        sah_plane = 1 + ((n_triangles_left * sleft) / sparent)
                        sah_planes[i] = sah_plane
                    else:
                        aabbleft = AABB(triangles_left.min(axis=1).min(axis=0), triangles_left.max(axis=1).max(axis=0))
                        aabbright = AABB(triangles_right.min(axis=1).min(axis=0), triangles_right.max(axis=1).max(axis=0))
                        diagleft = np.abs(aabbleft.topright - aabbleft.botleft)
                        diagright = np.abs(aabbright.topright - aabbright.botleft)
                        sleft = diagleft[0]*diagleft[1]+diagleft[1]*diagleft[2]+diagleft[0]*diagleft[2]
                        sright = diagright[0]*diagright[1]+diagright[1]*diagright[2]+diagright[0]*diagright[2]
                        sah_plane = 1+((n_triangles_left*sleft+n_triangles_right*sright)/sparent)
                        sah_planes[i] = sah_plane
                if (sah_planes.min()) < sah:
                    sah = sah_planes.min()
                    plane_sah = plane
                    axis_sah = axis
            sorttris = triangles[np.argsort(centroids[:, axis_sah])]
            if (len(sorttris[centroids[:, axis_sah] < plane_sah]))==0 or (len(sorttris[centroids[:, axis_sah] >= plane_sah]))==0:
                return BVHNode(aabb, triangles, None, None)

            node1 = build_bvh(sorttris[centroids[:, axis_sah] < plane_sah], mode='sah1')
            node2 = build_bvh(sorttris[centroids[:, axis_sah] >= plane_sah], mode='sah1')
            return BVHNode(aabb, np.zeros((0, 3, 3), dtype=np.float32), node1, node2)


def bvh_serializer(root):
    queue = deque()
    queue.append(root)
    nodelist = []

    while len(queue) > 0:
        front = queue.popleft()
        nodelist.append(front)
        if front.child1 is not None:
            queue.append(front.child1)
        if front.child2 is not None:
            queue.append(front.child2)

    for node in nodelist:
        node.parent = None

    ## generate serial node lists and triangle list
    serialnodelist = []
    serialtrilist = np.zeros((0, 3, 3), dtype=np.float32)
    serialnormallist = np.zeros((0, 3, 3), dtype=np.float32)
    serialtriindlist = np.zeros((0,), dtype=np.int32)
    for i, node in enumerate(nodelist):
        serialnodelist.append(SerialBVHNode(node.aabb,
                                            serialtrilist.shape[0],
                                            node.triangles.shape[0]))

        if node.child1 is not None:
            node.child1.parent = serialnodelist[-1]
        if node.child2 is not None:
            node.child2.parent = serialnodelist[-1]

        serialtrilist = np.vstack((serialtrilist, node.triangles))
        # serialnormallist = np.vstack((serialnormallist, node.vernormals))
        # serialtriindlist = np.hstack((serialtriindlist, node.triindices))

        ## find children by looking up parent
    for i, node, snode in zip(range(len(nodelist)), nodelist, serialnodelist):
        if (node.parent is not None):

            if (node.parent.child1 == -1):
                node.parent.child1 = i
                snode.parent = i
                continue

            if (node.parent.child2 == -1):
                node.parent.child2 = i
                snode.parent = i
                continue
    aabblist = np.zeros(((len(serialnodelist)), 2, 3), dtype=np.float32)
    ntrilist = np.zeros((len(serialnodelist)), dtype=np.int32)
    parentslist = np.zeros((len(serialnodelist)), dtype=np.int32)
    tristartlist = np.zeros((len(serialnodelist)), dtype=np.int32)
    childlist = np.zeros((len(serialnodelist), 2), dtype=np.int32)
    for i, node in enumerate(serialnodelist):
        tristartlist[i] = node.tristart
        ntrilist[i] = node.ntris
        aabblist[i] = [node.aabb.botleft, node.aabb.topright]
        childlist[i] = [node.child1, node.child2]
    for i in range(1, len(serialnodelist)):
        parentslist[i] = np.where(childlist == i)[0][0]
    parentslist[0] = -1
    # for i, node in enumerate(nodelist):
    #     print(node.parent)
    # for i, node in enumerate(nodelist):
    #     parentslist[i] = node.child1.parent
    # print('Дерево построено')
    return aabblist, serialtrilist, tristartlist, ntrilist, childlist, parentslist

def show_trees(childlist, aabblist):
    nodelist = [0]
    oldnodeslist = []
    i = 0
    # print(nodelist)
    # print(type(nodelist))
    level_trees(nodelist, oldnodeslist, childlist, aabblist, i)
    # vpl.show()

def show_plot_vpl(triangles, i):
    fig = vpl.figure()
    # vpl.figure()
    vpl.mesh_plot(triangles, fig=fig)
    vpl.save_fig(f'pic2/fig{i}.png')
    vpl.show(fig)
    # vpl.close()

def level_trees(nodeslist,oldnodeslist, childlist, aabblist, i):
    triangles = np.array([[[0,0,0], [0,0,0], [0,0,0]]], dtype=np.float32)
    nodeslist1 = []
    for node in nodeslist:
        lb = aabblist[node][0]
        rt = aabblist[node][1]
        aabb = AABB(lb, rt)
        aabb.create_box()
        triangles = np.vstack((triangles, aabb.triangles_aabb))
        if childlist[node][0] != -1:
            nodeslist1.append(childlist[node][0])
        if childlist[node][1] != -1:
            nodeslist1.append(childlist[node][1])
        if childlist[node][0] == -1 and childlist[node][1] == -1:
            oldnodeslist.append(node)
    for node in oldnodeslist:
        lb = aabblist[node][0]
        rt = aabblist[node][1]
        aabb = AABB(lb, rt)
        aabb.create_box()
        triangles = np.vstack((triangles, aabb.triangles_aabb))
    print(nodeslist)
    print(oldnodeslist)
    nodeslist.clear()
    nodeslist = nodeslist1.copy()
    show_plot_vpl(triangles[1:], i)
    i+=1
    if len(nodeslist) != 0:
        print(1)
        return level_trees(nodeslist, oldnodeslist, childlist, aabblist, i)


def main():
#    scene = Scene('uaz469.stl')
    triangles = scene.triangles_data
    print(triangles.shape[0])
    root = build_bvh(triangles, mode='median')
    aabblist, trilist, tristartlist, ntrilist, childlist, parents_list = bvh_serializer(root)
    print(len(childlist))
    show_trees(childlist, aabblist)


if __name__ == '__main__':
    main()
